import numpy as np

# This class implements a writer for a restricted form of the DIMACS graph format.
# This implementation assumes indices are base-0.  It supports my "extended" 
# weighted format.  It assumes the provided adjacency matrix is in numpy format.

# The rows of the adjacency matrix correspond to variables, and the columns give
# the parents for that variable.

class DimacsGraphWriter:

    def __init__(self):
        self.stream = None
        self.graph = None

    def write(self, graph, filename):
        # now output the cmi values
        out = open(filename, 'w')
        n = len(graph)

        out.write("p {} {}\n".format(n, graph.getEdgeCount()))

        for destination in range(n):
            for source in range(n):
                if graph.getEdge(source, destination) != 0:
                    if graph.weighted:
                        out.write("e {} {} {}\n".format(source, destination, graph.getEdge(source, destination)))
                    else:
                        out.write("e {} {}\n".format(source, destination))

        out.close()

    def write_adjacency_matrix(self, graph, filename, presence_only=True, sep='\t'):
        """ This method writes the adjacency matrix of the provided graph to
            the file. Optionally, the values of the adjacency matrix (i.e.,
            weights of the edges) can be written. By default, only "0" and "1"
            are written, which indicate the presence/absence of an edge. The
            file is tab-delimited, by default.

            The ROWS correspond to nodes. The COLUMNS correspond to parents.

            Args:
                graph (DimacsGraph) : the graph

                filename (string) : the path to the output file

                presence_only (bool) : whether to write only the presence of
                    an edge or its weight

                sep (string) : the separator for the file
        """

        if presence_only:
            graph = graph.toUnweighted()

        adjacency_matrix = graph.adjacencyMatrix

        with open(filename, 'w') as out:
            for row in adjacency_matrix:
                line = sep.join(str(v) for v in row)
                out.write(line)
                out.write("\n")
